import os
import pathlib

import pandas as pd

root_dir = pathlib.PurePath(
    os.path.dirname(os.path.realpath(__file__))
).parent.parent.parent
project_dir = os.path.join(root_dir, "resources", "sample_project")
qa_data = pd.read_parquet(os.path.join(project_dir, "data", "qa.parquet"))
corpus_data = pd.read_parquet(os.path.join(project_dir, "data", "corpus.parquet"))
doc_id_list = corpus_data["doc_id"].tolist()
previous_result = pd.DataFrame(
    {
        "qid": [1, 2],
        "query": [
            "What is the capital of France?",
            "How many members are in Newjeans?",
        ],
        "retrieved_ids": [
            [doc_id_list[1], doc_id_list[3]],
            [doc_id_list[5], doc_id_list[7]],
        ],
        "retrieved_contents": [
            [
                "Paris is the capital of France.",
                "Paris is one of the capital from France. Isn't it?",
            ],
            ["Newjeans has 5 members.", "Danielle is one of the members of Newjeans."],
        ],
        "retrieve_scores": [[0.8, 0.1], [0.2, 0.1]],
        "retrieval_gt": [1, 1],
        "generation_gt": ["answer", "answer"],
        "retrieval_f1": [0.4, 1.0],
        "retrieval_recall": [1.0, 0.3],
    }
)
ids_list = [[doc_id_list[1]], [doc_id_list[3]], [doc_id_list[0], doc_id_list[29]]]
